import torch
import torch.nn as nn
from torchvision import  models

def to_onnx(model, c, w, h, onnx_name):
    dummy_input = torch.randn(1, c, w, h, device='cpu')
    torch.onnx.export(model, dummy_input, onnx_name, verbose=True)


if __name__ == '__main__':
    model = models.resnet18(pretrained=False)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 1824)
    root = r'D:\proj\button\new_button\best_model_by_138.pth'
    xxx = torch.load(root, map_location='cpu')
    model.load_state_dict({k.replace("module.", ""):v for k, v in xxx.items()})
    # model.load_state_dict(torch.load(root, map_location='cpu'),False)
    model.eval()
    print('load model')
    to_onnx(model, 3, 160, 160, 'params.onnx')
